#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
MySQL MCP HTTP Server
基于HTTP的MySQL MCP服务器，接收JSON-RPC请求
"""

import json
import logging
import traceback
from typing import Any, Dict, List, Optional
from contextlib import asynccontextmanager

import uvicorn
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import pymysql
from pymysql.cursors import DictCursor

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 数据库配置
DB_CONFIG = {
    'host': 'localhost',
    'user': 'root',
    'password': '123456',
    'database': 'school',
    'charset': 'utf8mb4',
    'cursorclass': DictCursor,
    'autocommit': True,
    'auth_plugin': 'mysql_native_password'  # 使用传统认证方式
}

class JSONRPCRequest(BaseModel):
    jsonrpc: str = "2.0"
    id: Optional[str] = None
    method: str
    params: Optional[Dict[str, Any]] = None

class JSONRPCResponse(BaseModel):
    jsonrpc: str = "2.0"
    id: Optional[str] = None
    result: Optional[Dict[str, Any]] = None
    error: Optional[Dict[str, Any]] = None

class MCPTool(BaseModel):
    name: str
    description: str
    inputSchema: Dict[str, Any]

class MCPTools(BaseModel):
    tools: List[MCPTool]

# 全局变量
app = FastAPI(title="MySQL MCP HTTP Server", version="1.0.0")

# 添加CORS中间件
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

def get_db_connection():
    """获取数据库连接"""
    # 尝试不同的认证方式
    auth_methods = [
        {'auth_plugin': 'mysql_native_password'},
        {'auth_plugin': 'caching_sha2_password'},
        {}  # 默认方式
    ]
    
    for auth_config in auth_methods:
        try:
            config = DB_CONFIG.copy()
            config.update(auth_config)
            connection = pymysql.connect(**config)
            logger.info(f"数据库连接成功，使用认证方式: {auth_config.get('auth_plugin', 'default')}")
            return connection
        except Exception as e:
            logger.warning(f"认证方式 {auth_config.get('auth_plugin', 'default')} 失败: {e}")
            continue
    
    # 所有认证方式都失败
    logger.error("所有数据库认证方式都失败")
    raise HTTPException(status_code=500, detail="数据库连接失败: 无法建立连接")

def execute_query(query: str, params: Optional[tuple] = None) -> Dict[str, Any]:
    """执行SQL查询"""
    connection = None
    try:
        connection = get_db_connection()
        with connection.cursor() as cursor:
            cursor.execute(query, params)
            
            # 判断查询类型
            query_upper = query.strip().upper()
            if query_upper.startswith('SELECT') or query_upper.startswith('SHOW') or query_upper.startswith('DESCRIBE'):
                # 查询操作
                results = cursor.fetchall()
                return {
                    "success": True,
                    "data": results,
                    "row_count": len(results)
                }
            else:
                # 修改操作
                connection.commit()
                affected_rows = cursor.rowcount
                return {
                    "success": True,
                    "affected_rows": affected_rows,
                    "message": "操作成功"
                }
    except Exception as e:
        logger.error(f"查询执行失败: {e}")
        return {
            "success": False,
            "error": str(e)
        }
    finally:
        if connection:
            connection.close()

def get_metadata() -> Dict[str, Any]:
    """获取数据库元数据"""
    try:
        connection = get_db_connection()
        metadata = {
            "tables": [],
            "database_name": DB_CONFIG['database']
        }
        
        with connection.cursor() as cursor:
            # 获取所有表
            cursor.execute("SHOW TABLES")
            tables = cursor.fetchall()
            
            for table in tables:
                table_name = list(table.values())[0]
                
                # 获取表结构
                cursor.execute(f"DESCRIBE `{table_name}`")
                columns = cursor.fetchall()
                
                # 获取记录数
                cursor.execute(f"SELECT COUNT(*) as count FROM `{table_name}`")
                count_result = cursor.fetchone()
                record_count = count_result['count'] if count_result else 0
                
                # 构建表信息
                table_info = {
                    "name": table_name,
                    "columns": [],
                    "record_count": record_count
                }
                
                for col in columns:
                    column_info = {
                        "name": col['Field'],
                        "type": col['Type'],
                        "null": col['Null'] == 'YES',
                        "key": col['Key'],
                        "default": col['Default'],
                        "extra": col['Extra']
                    }
                    table_info["columns"].append(column_info)
                
                metadata["tables"].append(table_info)
        
        connection.close()
        return metadata
        
    except Exception as e:
        logger.error(f"获取元数据失败: {e}")
        raise HTTPException(status_code=500, detail=f"获取元数据失败: {e}")

    """处理工具列表请求"""
    tools = [
        {
            "name": "get_metadata",
            "description": "获取数据库元数据，包括表结构和记录数",
            "inputSchema": {
                "type": "object",
                "properties": {
                    "protocolVersion": {"type": "string"},
                    "capabilities": {"type": "object"},
                    "clientInfo": {"type": "object"}
                },
                "required": ["protocolVersion", "capabilities", "clientInfo"]
            }
        },
        {
            "name": "execute_query",
            "description": "执行SQL查询",
            "inputSchema": {
                "type": "object",
                "properties": {
                    "query": {"type": "string", "description": "要执行的SQL查询"}
                },
                "required": ["query"]
            }
        }
    ]
    
    return {"tools": tools}

def handle_tools_call(name: str, arguments: Dict[str, Any]) -> Dict[str, Any]:
    """处理工具调用请求"""
    if name == "get_metadata":
        metadata = get_metadata()
        return {
            "content": [
                {
                    "type": "text",
                    "text": json.dumps(metadata, ensure_ascii=False, indent=2)
                }
            ]
        }
    elif name == "execute_query":
        query = arguments.get("query", "")
        if not query:
            raise HTTPException(status_code=400, detail="查询语句不能为空")
        
        result = execute_query(query)
        return {
            "content": [
                {
                    "type": "text",
                    "text": json.dumps(result, ensure_ascii=False, indent=2)
                }
            ]
        }
    else:
        raise HTTPException(status_code=400, detail=f"未知的工具: {name}")

@app.get("/health")
async def health_check():
    """健康检查"""
    try:
        # 检查数据库连接
        connection = get_db_connection()
        connection.close()
        
        return {
            "status": "healthy",
            "database": "connected",
            "server": "running"
        }
    except Exception as e:
        return {
            "status": "unhealthy",
            "database": "disconnected",
            "server": "running",
            "error": str(e)
        }

@app.post("/rpc")
async def handle_rpc_post(request: JSONRPCRequest):
    """处理JSON-RPC请求"""
    try:
        logger.info(f"收到RPC请求: {request.method}")
        
        if request.method == "tools/call":
            if not request.params:
                raise HTTPException(status_code=400, detail="缺少参数")
            
            name = request.params.get("name")
            arguments = request.params.get("arguments", {})
            
            if not name:
                raise HTTPException(status_code=400, detail="缺少工具名称")
            
            result = handle_tools_call(name, arguments)
        else:
            raise HTTPException(status_code=400, detail=f"未知的方法: {request.method}")
        
        return JSONRPCResponse(
            id=request.id,
            result=result
        )
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"RPC处理错误: {e}")
        logger.error(traceback.format_exc())
        
        return JSONRPCResponse(
            id=request.id,
            error={
                "code": -32603,
                "message": "Internal error",
                "data": str(e)
            }
        )

@app.get("/rpc")
async def handle_rpc_get(method: str = "tools/call", name: str = "get_metadata", arguments: str = "{}"):
    """处理GET请求的JSON-RPC调用"""
    try:
        logger.info(f"收到GET RPC请求: {method}")
        
        # 解析arguments参数
        try:
            import json
            parsed_arguments = json.loads(arguments) if arguments else {}
        except json.JSONDecodeError:
            parsed_arguments = {}
        
        if method == "tools/call":
            result = handle_tools_call(name, parsed_arguments)
        else:
            raise HTTPException(status_code=400, detail=f"未知的方法: {method}")
        
        return JSONRPCResponse(
            id="1",
            result=result
        )
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"GET RPC处理错误: {e}")
        logger.error(traceback.format_exc())
        
        return JSONRPCResponse(
            id="1",
            error={
                "code": -32603,
                "message": "Internal error",
                "data": str(e)
            }
        )

@app.get("/")
async def root():
    """根路径"""
    return {
        "message": "MySQL MCP HTTP Server",
        "version": "1.0.0",
        "endpoints": {
            "health": "/health",
            "rpc": "/rpc",
            "docs": "/docs"
        }
    }

if __name__ == "__main__":
    print("启动MySQL MCP HTTP服务器...")
    print("访问 http://localhost:8081/docs 查看API文档")
    print("RPC端点: http://localhost:8081/rpc")
    print("按 Ctrl+C 停止服务器")
    
    uvicorn.run(
        "mysql_http_server:app",
        host="0.0.0.0",
        port=8081,
        reload=False,
        log_level="info"
    )